#include <iostream>
#include <sstream>

#include <vector>
#include <map>
#include <set>
#include <memory>
#include <list>

#include <cassert>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <cluster.h>
#include <segment.h>
#include <matrix.h>
#include <random.h>

#include <cassert>

#include <lemon/grid_graph.h>

using namespace std;
using namespace sds;
using namespace lemon;

typedef SDataView<int> SegmentView;
typedef SDataView<Point> PointView;
typedef SDataView<unsigned char> CharView;

void Read(const string& filename, GridGraph& graph,
	  GridGraph::EdgeMap<double>& weight, vector<SInteger>& size) {
  CharView source[CH];
  SFile file(filename);
  for (int k = 0; k < CH; ++k) {
    source[k] = file.getDataset(k).getDataView<unsigned char>();
  }
  
  size = source[0].getSize();
 
  Cluster dist;
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      dist += point;
    }
  }

  Matrix cm(CH, CH);
  for (int i = 0; i < CH; ++i) {
    for (int j = 0; j < CH; ++j) {
      cm(i, j) = dist.covariance(i, j);
    }
  }
  Matrix cmi = cm.inverse();

  graph.resize(size[0], size[1]);

  for (GridGraph::EdgeIt it(graph); it != INVALID; ++it) {
    Vector dv(CH);
    for (int k = 0; k < CH; ++k) {
      dv(k) = (double)source[k][graph.col(graph.u(it))][graph.row(graph.u(it))]-
	(double)source[k][graph.col(graph.v(it))][graph.row(graph.v(it))];
    }
    //    std::cerr << -1 * dv * cmi * dv / CH << ' ';
    weight[it] = exp(-1 * dv * cmi * dv);
    //    std::cerr << weight[it] << std::endl;
  }
}

class Splitter {
  const GridGraph& _graph;
  const GridGraph::EdgeMap<double>& _weight;

  GridGraph::NodeMap<int>& _segments;

  GridGraph::NodeMap<int> _index;

  std::vector<int> _size;

  std::vector<GridGraph::Node> _coord;
  std::vector<double> _sqdiag;
  double _sum;

  std::vector<double> _value;
  std::vector<int> _row;
  std::vector<int> _scol;
  std::vector<bool> _red;

  int _num;

  Random _rnd;

public:

  Splitter(const GridGraph& graph, 
	   const GridGraph::EdgeMap<double>& weight, 
	   GridGraph::NodeMap<int>& segments) 
    : _graph(graph), _weight(weight), _segments(segments), 
      _index(graph, -1), _num(0)
  {}
  
  // (x,y) is always the left-top corner of a segment
  bool split(GridGraph::Node p, std::vector<GridGraph::Node>& nps) {
    if (singular(p)) return false;
    createGraph(p);
    bool success = extractRegions(nps);
    clearGraph();
    return success;
  }

private:

  bool singular(GridGraph::Node p) {
    static const int dn = 4;
    
    static const int dv[] = {-1,  0,  1,  0 };
    static const int dh[] = { 0,  1,  0, -1 };
    
    int seg = _segments[p];    

    for (GridGraph::OutArcIt a(_graph, p); a != INVALID; ++a) {
      GridGraph::Node np = _graph.target(a);
      if (_segments[np] == seg) {
	return false;
      }
    }
    return true;
  }

  void createGraph(GridGraph::Node p) {
    static const int dn = 4;
    
    int seg = _segments[p];

    _index[p] = _coord.size();
    _coord.push_back(p);

    int cur = 0;
    while (cur < static_cast<int>(_coord.size())) {
      GridGraph::Node pp = _coord[cur++];

      for (GridGraph::OutArcIt a(_graph, pp); a != INVALID; ++a) {
	GridGraph::Node np = _graph.target(a);

	if (_segments[np] == seg && _index[np] == -1) {
	  _index[np] = _coord.size();
	  _coord.push_back(np);
	}
      }
    }

    _sum = 0.0;

    cur = 0;
    while (cur < static_cast<int>(_coord.size())) {
      int id = cur;
      GridGraph::Node pp = _coord[cur++];

      _scol.push_back(_value.size());
      _red.push_back((_graph.col(pp) - _graph.row(pp)) % 2 == 0);

      double sum = 0.0;

      for (GridGraph::OutArcIt a(_graph, pp); a != INVALID; ++a) {
	GridGraph::Node np = _graph.target(a);
	if (_segments[np] == seg) {
	  double dist = _weight[a];

	  _row.push_back(_index[np]);
	  _value.push_back(-dist);
	  sum += dist;
	}
      }
      _row.push_back(id);
      _value.push_back(sum);
      _sqdiag.push_back(std::sqrt(sum));
      _sum += sum;
    }
    _scol.push_back(_value.size());

    for (int i = 0; i < static_cast<int>(_coord.size()); ++i) {
      for (int j = _scol[i]; j < _scol[i + 1]; ++j) {
	//	std::cerr << _value[j] << ' ';
	_value[j] /= _sqdiag[i] * _sqdiag[_row[j]];
	//	std::cerr << _value[j] << std::endl;
      }
    }
  }

  bool extractRegions(std::vector<GridGraph::Node>& nps) {
    std::vector<double> res(_coord.size());
    double cross;

    for (int i = 0; i < static_cast<int>(_coord.size()); ++i) {
      res[i] = _rnd.stdnorm();
    }

    for (int k = 0; k < 15000; ++k) {
      // matrix multiplication
      std::vector<double> old(res);
      for (int i = 0; i < static_cast<int>(_coord.size()); ++i) {
	res[i] = 0.0;
	for (int j = _scol[i]; j < _scol[i + 1]; ++j) {
	  res[i] += old[_row[j]] * _value[j];
	}	
      }

      // project vector to be transversal to the eigenvector with
      // largest eigenvalue
      cross = 0.0;
      for (int i = 0; i < static_cast<int>(_coord.size()); ++i) {
      	cross += res[i] * (_red[i] ? 1 : -1) * _sqdiag[i];	
      }
      cross /= _sum;
      for (int i = 0; i < static_cast<int>(_coord.size()); ++i) {
      	res[i] -= cross * (_red[i] ? 1 : -1) * _sqdiag[i];
      }

      // normalize vector
      cross = 0.0; 
      for (int i = 0; i < static_cast<int>(_coord.size()); ++i) {
      	cross += res[i] * res[i];
      }
      for (int i = 0; i < static_cast<int>(_coord.size()); ++i) {
      	res[i] /= std::sqrt(cross);
      }
    }

    std::vector<std::pair<double, int> > pairs;
    for (int i = 0; i < static_cast<int>(_coord.size()); ++i) {
      pairs.push_back(std::make_pair(res[i] * (_red[i] ? 1 : -1), i));
    }

    std::sort(pairs.begin(), pairs.end());

    std::vector<int> order(pairs.size());
    for (int i = 0; i < static_cast<int>(pairs.size()); ++i) {
      order[pairs[i].second] = i;      
    }

    double mnc = std::numeric_limits<double>::max();
    int mnci = -1;

    double cut = 0.0, d1 = 0.0, d2 = _sum;
    for (int i = 0; i < static_cast<int>(pairs.size()); ++i) {
      int id = pairs[i].second;
      for (GridGraph::OutArcIt a(_graph, _coord[id]); a != INVALID; ++a) {
	if (_segments[_coord[id]] != _segments[_graph.target(a)]) continue;
	int jd = _index[_graph.target(a)];
	int j = order[jd];
	if (i < j) cut += _weight[a];
	else if (i > j) cut -= _weight[a];
	d1 += _weight[a];
	d2 -= _weight[a];
      }
      if (d1 > 0 && d2 > 0 && cut > 0 && i != 0 && i + 1 != pairs.size() && 
	  cut / d1 + cut / d2 < mnc) {
	if (pairs.size() < 3000 || (i >= pairs.size() / 100 &&
				    i < 99 * pairs.size() / 100)) {
	  mnc = cut / d1 + cut / d2;
	  mnci = i;
	}
      }
      // std::cerr << d1 << ' ' << d2 << ' ' << cut << ' '
      // 		<< cut / d1 + cut / d2 << ' ' << mnc << std::endl;
    }

    std::cerr << mnc << ' ' << mnci + 1 << ' ' 
	      << pairs.size() - mnci - 1 << std::endl;
    if (mnc > 0.015) return false;

    std::vector<char> state(order.size(), 0);

    ++_num;

    int seg = _segments[_coord[0]];
    
    std::vector<GridGraph::Node> sources;
    sources.push_back(_coord[0]);
    state[_index[_coord[0]]] = 1;
    
    while (!sources.empty()) {
      GridGraph::Node p = sources.back();
      sources.pop_back();
      if (state[_index[p]] == 2) continue;
      nps.push_back(p);

      std::vector<GridGraph::Node> stack;
      stack.push_back(p);
      state[_index[p]] = 2;
      ++_num;
      
      while (!stack.empty()) {
  	GridGraph::Node pp = stack.back();
  	stack.pop_back();
  	_segments[pp] = _num;
	
  	for (GridGraph::OutArcIt a(_graph, pp); a != INVALID; ++a) {
	  if (_segments[_graph.target(a)] == seg) {
	    GridGraph::Node np = _graph.target(a);
	    if ((order[_index[pp]] > mnci) != (order[_index[np]] > mnci)) {
	      if (state[_index[np]] == 0) {
		sources.push_back(np);
		state[_index[np]] = 1;
	      }
	    } else {
	      if (state[_index[np]] == 0 || state[_index[np]] == 1) {
		stack.push_back(np);
		state[_index[np]] = 2;
	      }
	    }
  	  }
  	}
      }
    }
    return true;
  }

  void clearGraph() {
    int cur = 0;
    while (cur < static_cast<int>(_coord.size())) {
      GridGraph::Node pp = _coord[cur++];
      _index[pp] = -1;
    }
    _value.clear();
    _row.clear();
    _scol.clear();
    _coord.clear();
    _red.clear();
    _sqdiag.clear();
  }
};

int main(int argc, const char* argv[]) {

  PointView points;
  std::vector<SInteger> size;

  GridGraph graph(0, 0);
  GridGraph::EdgeMap<double> weight(graph);
  GridGraph::NodeMap<int> segments(graph, 0);

  Read(argv[1], graph, weight, size);

  for (GridGraph::NodeIt n(graph); n != INVALID; ++n) {
    segments[n] = 0;
  }
  
  //  SegmentView segments = 
  //    SFile(argv[2]).getDataset(0).getDataView<int>();
  Splitter splitter(graph, weight, segments);
  std::list<GridGraph::Node> queue;
  queue.push_back(graph(0, 0));
  
  int k = 0;
  while (!queue.empty()) {
    std::cerr << ++k << std::endl;
    std::vector<GridGraph::Node> pns;
    splitter.split(queue.front(), pns);
    queue.pop_front();
    queue.insert(queue.end(), pns.begin(), pns.end());

    // if (k < 15 || k % 15 == 0 || queue.empty()) {
    //   SDataView<int> seg(size);
    //   for (int i = 0; i < size[0]; ++i) {
    // 	for (int j = 0; j < size[1]; ++j) {
    // 	  seg[i][j] = segments[graph(i, j)];
    // 	}
    //   }
    //   SDataView<unsigned char> trg(size);
      
    //   WriteSegments(seg, trg);
    
    //   // std::ostringstream oss;
    //   // oss << "normcutsegmentation_" << k << ".hdf";
    
    //   SFile trgfile(oss.str(), SFile::CREATE);
    //   trgfile.createDataset("segmented", sds::UINT8, size).
    // 	setDataView<unsigned char>(trg);
    // }
    
  }
  
  SDataView<int> seg(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      seg[i][j] = segments[graph(i, j)];
    }
  }

  SDataView<unsigned char> trg(size);
  
  WriteSegments(seg, trg);
  
  // std::ostringstream oss;
  // oss << "normcutsegmentation_" << k << ".hdf";
    
  SFile trgfile(argv[2], SFile::CREATE);
  trgfile.createDataset("segmented", sds::UINT8, size).
    setDataView<unsigned char>(trg);
  
  return 0;
}

